#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简化版数据获取器 - 专门为中长期投资分析优化
只保留必要的功能，提高稳定性和速度
"""

import akshare as ak
import pandas as pd
import numpy as np
import time
from datetime import datetime, timedelta
import logging
import warnings
import os

# 禁用代理
os.environ['HTTP_PROXY'] = ''
os.environ['HTTPS_PROXY'] = ''
os.environ['http_proxy'] = ''
os.environ['https_proxy'] = ''

warnings.filterwarnings('ignore')

logger = logging.getLogger(__name__)

class SimplifiedDataFetcher:
    """简化版股票数据获取器"""
    
    def __init__(self):
        self.retry_count = 3
        self.delay = 1  # 请求间隔
        
    def get_stock_list(self):
        """获取A股股票列表"""
        try:
            # 使用正确的API获取股票列表
            stock_list = ak.stock_info_a_code_name()
            
            if stock_list.empty:
                return pd.DataFrame()
            
            # 获取实时行情数据来补充价格和市值信息
            try:
                spot_data = ak.stock_zh_a_spot_em()
                # 合并数据
                merged_data = stock_list.merge(
                    spot_data[['代码', '最新价', '总市值']], 
                    left_on='code', 
                    right_on='代码', 
                    how='left'
                )
                
                # 简化字段
                simplified_list = pd.DataFrame({
                    'code': merged_data['code'],
                    'name': merged_data['name'],
                    'price': pd.to_numeric(merged_data['最新价'], errors='coerce').fillna(0),
                    'market_cap': pd.to_numeric(merged_data['总市值'], errors='coerce').fillna(0) / 100000000  # 转换为亿元
                })
                
            except Exception as e:
                logger.warning(f"获取实时行情失败，使用基础股票列表: {e}")
                # 如果获取实时数据失败，只返回股票代码和名称
                simplified_list = pd.DataFrame({
                    'code': stock_list['code'],
                    'name': stock_list['name'],
                    'price': 10.0,  # 默认价格
                    'market_cap': 100.0  # 默认市值100亿
                })
            
            # 过滤掉无效数据
            simplified_list = simplified_list[
                (simplified_list['price'] > 0) & 
                (simplified_list['market_cap'] > 0)
            ].copy()
            
            logger.info(f"获取到 {len(simplified_list)} 只A股数据")
            return simplified_list
            
        except Exception as e:
            logger.error(f"获取股票列表失败: {e}")
            return pd.DataFrame()
    
    def get_stock_data(self, stock_code: str, days: int = 120):
        """获取股票历史数据"""
        for attempt in range(self.retry_count):
            try:
                time.sleep(self.delay)
                
                # 计算日期范围
                end_date = datetime.now()
                start_date = end_date - timedelta(days=days + 30)  # 多获取一些数据以防节假日
                
                # 获取历史数据
                hist_data = ak.stock_zh_a_hist(
                    symbol=stock_code,
                    period="daily",
                    start_date=start_date.strftime('%Y%m%d'),
                    end_date=end_date.strftime('%Y%m%d'),
                    adjust="qfq"  # 前复权
                )
                
                if hist_data.empty:
                    continue
                
                # 标准化列名
                hist_data.columns = ['date', 'open', 'close', 'high', 'low', 'volume', 'amount', 'amplitude', 'change_pct', 'change_amount', 'turnover']
                
                # 转换数据类型
                hist_data['date'] = pd.to_datetime(hist_data['date'])
                numeric_columns = ['open', 'close', 'high', 'low', 'volume', 'amount']
                for col in numeric_columns:
                    hist_data[col] = pd.to_numeric(hist_data[col], errors='coerce')
                
                # 设置索引
                hist_data.set_index('date', inplace=True)
                
                # 排序并获取最近的数据
                hist_data = hist_data.sort_index().tail(days)
                
                logger.info(f"获取股票 {stock_code} 数据成功: {len(hist_data)} 条记录")
                return hist_data
                
            except Exception as e:
                logger.warning(f"获取股票 {stock_code} 数据失败 (尝试 {attempt + 1}/{self.retry_count}): {e}")
                if attempt < self.retry_count - 1:
                    time.sleep(2 ** attempt)  # 指数退避
                    
        logger.error(f"获取股票 {stock_code} 数据最终失败")
        return pd.DataFrame()
    
    def get_stock_basic_info(self, stock_code: str):
        """获取股票基本信息"""
        try:
            time.sleep(self.delay)
            
            # 获取实时行情
            spot_data = ak.stock_zh_a_spot_em()
            stock_info = spot_data[spot_data['代码'] == stock_code]
            
            if stock_info.empty:
                return {}
            
            stock_info = stock_info.iloc[0]
            
            basic_info = {
                'code': stock_code,
                'name': stock_info['名称'],
                'close': float(stock_info['最新价']),
                'market_cap': float(stock_info['总市值']) / 100000000,  # 转换为亿元
                'pe_ratio': float(stock_info.get('市盈率-动态', 0)) if pd.notna(stock_info.get('市盈率-动态', 0)) else 0,
                'pb_ratio': float(stock_info.get('市净率', 0)) if pd.notna(stock_info.get('市净率', 0)) else 0,
                'turnover_rate': float(stock_info.get('换手率', 0)) if pd.notna(stock_info.get('换手率', 0)) else 0
            }
            
            return basic_info
            
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 基本信息失败: {e}")
            return {}
    
    def get_dividend_info(self, stock_code: str):
        """获取分红信息"""
        try:
            time.sleep(self.delay)
            
            # 获取分红数据
            dividend_data = ak.stock_history_dividend_detail(symbol=stock_code)
            
            if dividend_data.empty:
                return pd.DataFrame()
            
            # 筛选最近5年的分红记录
            current_year = datetime.now().year
            recent_dividends = []
            
            for _, row in dividend_data.iterrows():
                try:
                    announce_date = pd.to_datetime(row['公告日期'])
                    if announce_date.year >= current_year - 5:
                        dividend_amount = float(row.get('派息', 0))
                        if dividend_amount > 0:
                            recent_dividends.append({
                                'announce_date': announce_date,
                                'dividend_amount': dividend_amount,
                                'ex_dividend_date': row.get('除权除息日', ''),
                                'status': row.get('进度', '')
                            })
                except:
                    continue
            
            if recent_dividends:
                dividend_df = pd.DataFrame(recent_dividends)
                dividend_df = dividend_df.sort_values('announce_date', ascending=False)
                logger.info(f"获取股票 {stock_code} 分红信息: {len(dividend_df)} 条记录")
                return dividend_df
            else:
                return pd.DataFrame()
                
        except Exception as e:
            logger.warning(f"获取股票 {stock_code} 分红信息失败: {e}")
            return pd.DataFrame()
    
    def get_financial_data(self, stock_code: str):
        """获取财务数据"""
        try:
            time.sleep(self.delay)
            
            # 获取主要财务指标
            financial_data = ak.stock_financial_em(symbol=stock_code)
            
            if financial_data.empty:
                return {}
            
            latest_data = financial_data.iloc[-1]  # 最新一期数据
            
            financial_info = {
                'roe': float(latest_data.get('净资产收益率', 0)) if pd.notna(latest_data.get('净资产收益率', 0)) else 0,
                'gross_margin': float(latest_data.get('销售毛利率', 0)) if pd.notna(latest_data.get('销售毛利率', 0)) else 0,
                'debt_ratio': float(latest_data.get('资产负债率', 0)) if pd.notna(latest_data.get('资产负债率', 0)) else 0,
                'revenue_growth': float(latest_data.get('营业收入同比增长', 0)) if pd.notna(latest_data.get('营业收入同比增长', 0)) else 0,
                'net_profit_growth': float(latest_data.get('净利润同比增长', 0)) if pd.notna(latest_data.get('净利润同比增长', 0)) else 0
            }
            
            return financial_info
            
        except Exception as e:
            logger.warning(f"获取股票 {stock_code} 财务数据失败: {e}")
            return {}

# 为了兼容性，创建别名
StockDataFetcher = SimplifiedDataFetcher
